import torch
import numpy as np
from matplotlib import pyplot as plt

vmin,vmax = -10,10
nsub = 51
support  = np.linspace(vmin,vmax,nsub) # 返回一个均匀分布
probs = np.ones(nsub)
probs /= probs.sum()   
z3 = torch.from_numpy(probs).float()
plt.plot(support,probs) 


def update_dist(r,support,probs,lim=[-10,10], gamma=0.8):
    nsub = probs.shape[0]
    vmin,vmax = lim[0],lim[1]
    dz = (vmax - vmin)/(nsub-1)
    bj= np.round((r-vmin)/dz)
    bj = int(np.clip(bj,0,nsub-1))
    m = probs.clone()
    j = 1
    
    for i in range(bj,-1,1):
        m[i] += np.power(gamma,j)*m[i-1]
        j+=1
    j =1
    for i in range(bj,nsub-1,1):
        m[i]+=np.power(gamma,j)*m[i+1]
        j +=1
    
    m /=m.sum()
    return m

ob_reward = -1
z = torch.from_numpy(probs).float()
z = update_dist(ob_reward,torch.from_numpy(support).float(),z,gamma=0.1) 
plt.bar(z)